def add_obj(session, obj):
    """ Check if object primary key exists in db. If so,exit, else
    add.
    """
    from sqlalchemy import and_
    from sqlalchemy.orm import object_mapper
    mapper = object_mapper(obj)
    pid = mapper.primary_key_from_instance(obj)
    criterion = and_(*(col == val for col, val in zip(mapper.primary_key, mapper.primary_key_from_instance(obj))))
    if session.query(obj.__class__).filter(criterion).count() > 0:
        print "%s object with id %s is already in db."%(type(obj).__name__, pid)
        exit
    else:
        session.add(obj)

def get_obj(session, cls):
    query = session.query(cls)
    print "*** list of %s objects in class. ***"%cls.__name__
    for obj in query:
        print obj.__dict__
        print "%s is %s"%(cls.__name__, obj)
    print "*** end list of %s objects. ***"%cls.__name__

def get_db_type(dbstring):
   from utils import create_engine_wrapper
   db = create_engine_wrapper(dbstring)
   conn = db.connect()
   from sqlalchemy.sql import text
   q1 = text("SELECT current_database();")
   dbname = conn.execute(q1).fetchall()[0][0]
   q2 = text("SELECT oid, datname FROM pg_database where datname='%s';"%dbname)
   oid = conn.execute(q2).fetchall()[0][0]
   q3 = text("SELECT * from shobj_description(%s,'pg_database');"%oid)
   result = conn.execute(q3).fetchall()
   conn.close()
   db.dispose()
   return result[0][0]

def db_exists(dbengine, dbname):
    conn = dbengine.connect()
    result = conn.execute("select datname from pg_database where datname='%s'"%dbname).fetchall()
    conn.close()
    if len(result) > 0 and result[0][0]==dbname:
        return True
    else:
        return False

def dbuser_exists(dbengine, username):
    conn = dbengine.connect()
    result = conn.execute("select usename from pg_user where usename = '%s'"%username).fetchall()
    conn.close()
    if len(result) > 0 and result[0][0]==username:
        return True
    else:
        return False

def schema_exists(dbengine, schema, do_print=False):
    conn = dbengine.connect()
    result = conn.execute("""select schema_name from information_schema.schemata where schema_name='%s';"""%schema).fetchall()
    rv = len(result) > 0 and result[0][0]==schema
    if rv and do_print:
        print 'schema %s already exists'%schema
    return rv
    conn.close()

def tables_match_geno_single(conn, schema):
    from sqlalchemy.sql import text
    from utils import text_autocommit
    q = text_autocommit("select table_name from information_schema.tables where table_schema='%s' and table_name ~ '^geno$'; "%(schema))
    result = conn.execute(q).fetchall()
    if len(result) > 0:
        return True
    else:
        return False

def tables_match_geno_shard(conn, schema):
    from sqlalchemy.sql import text
    from utils import text_autocommit
    q = text_autocommit("select table_name from information_schema.tables where table_schema='%s' and table_name ~ '^geno[0-9]+$'; "%(schema))
    result = conn.execute(q).fetchall()
    if len(result) > 0:
        return True
    else:
        return False


def drop_schema(dbengine, schemaname):
    from utils import create_engine_wrapper
    conn = dbengine.connect()
    print "deleting schema %s"%schemaname
    conn.execute("commit")
    conn.execute("DROP SCHEMA IF EXISTS %s CASCADE"%schemaname)
    conn.execute("commit")
    conn.close()

def recreate_schema(dbengine, schemaname):
    from utils import create_engine_wrapper
    conn = dbengine.connect()
    print "recreating schema %s"%schemaname
    conn.execute("commit")
    conn.execute("drop schema %s cascade"%schemaname)
    conn.execute("commit")
    conn.execute("create schema %s"%schemaname)
    conn.close()

